#include <fltKernel.h>

#define NTSTRSAFE_NO_CB_FUNCTIONS
#include <ntstrsafe.h>
#include "../HyperPlatform/common.h"
#include "../HyperPlatform/kernel_stl.h"
#include "../../Shared/Protocol.h"
#include <set>
#include <vector>

#include "main.h"

PFLT_FILTER m_pFilterHandle = NULL;
PFLT_PORT 	m_pServerPort = NULL;
PFLT_PORT 	m_pClientPort = NULL;

extern CProcList *m_IgnoreProcList;
extern CFileList *m_IgnoreFileList;
extern CEventList *m_EventList;

NTSTATUS NtFileNameToDosFileName(IN PUNICODE_STRING NtFileName, OUT PUNICODE_STRING DosFileName);

EXTERN_C{

typedef struct
{
	UNICODE_STRING NtFilePath;
	UNICODE_STRING DosFilePath;
}flt_file_context_t;

NTSTATUS FLTAPI FsMessageNotifyCallback(
	IN PVOID PortCookie,
	IN PVOID InputBuffer OPTIONAL,
	IN ULONG InputBufferLength,
	OUT PVOID OutputBuffer OPTIONAL,
	IN ULONG OutputBufferLength,
	OUT PULONG ReturnOutputBufferLength);

VOID MessageSenderThread(IN PVOID pContext)
{
	UNREFERENCED_PARAMETER(pContext);

	while (!m_EventList->m_Stop)
	{
		svc_nop_data *msg = NULL;

		if (NT_SUCCESS(KeWaitForSingleObject(&m_EventList->m_MsgEvent, Executive, KernelMode, FALSE, 0)))
		{
			m_EventList->Lock();
			if (!m_EventList->m_List.empty())
			{
				msg = (svc_nop_data *)m_EventList->m_List.front();
				m_EventList->m_List.pop_front();

				if (!m_EventList->m_List.empty())
					KeSetEvent(&m_EventList->m_MsgEvent, IO_NO_INCREMENT, FALSE);
			}
			m_EventList->Unlock();

			if (msg)
			{
				if (m_pClientPort != NULL)
					FltSendMessage(m_pFilterHandle, &m_pClientPort, msg, msg->size, NULL, NULL, NULL);

				ExFreePoolWithTag(msg, 'TXSB');
			}
		}
	}

	PsTerminateSystemThread(STATUS_SUCCESS);
}

NTSTATUS FLTAPI FsConnectNotifyCallback(IN PFLT_PORT ClientPort, IN PVOID ServerPortCookie, IN PVOID ConnectionContext, IN ULONG SizeOfContext, OUT PVOID * ConnectionPortCookie)
{
	PAGED_CODE();

	UNREFERENCED_PARAMETER(ServerPortCookie);
	UNREFERENCED_PARAMETER(ConnectionContext);
	UNREFERENCED_PARAMETER(SizeOfContext);
	UNREFERENCED_PARAMETER(ConnectionPortCookie);

	if (SizeOfContext == sizeof(conn_context_data))
	{
		conn_context_data *data = (conn_context_data *)ConnectionContext;
		if (data->txsb == 'TXSB' && data->ver == 1)
		{
			m_pClientPort = ClientPort;

			return STATUS_SUCCESS;
		}
	}

	return STATUS_OBJECT_NAME_NOT_FOUND;
}

VOID FLTAPI FsDisconnectNotifyCallback(_In_opt_ PVOID ConnectionCookie)
{
	PAGED_CODE();

	UNREFERENCED_PARAMETER(ConnectionCookie);

	FltCloseClientPort(m_pFilterHandle, &m_pClientPort);
	m_pClientPort = NULL;
}

NTSTATUS FLTAPI FsUnload(__in FLT_FILTER_UNLOAD_FLAGS Flags)
{
	UNREFERENCED_PARAMETER(Flags);
	PAGED_CODE();

	HYPERPLATFORM_COMMON_DBG_BREAK();

	InterlockedExchange(&m_EventList->m_Stop, 1);	
	KeSetEvent(&m_EventList->m_MsgEvent, IO_NO_INCREMENT, FALSE);
	ZwWaitForSingleObject(m_EventList->m_hMsgThread, FALSE, NULL);
	m_EventList->m_hMsgThread = NULL;

	if (NULL != m_pServerPort) {
		FltCloseCommunicationPort(m_pServerPort);
		m_pServerPort = NULL;
	}
	if (NULL != m_pClientPort) {
		FltCloseClientPort(m_pFilterHandle, &m_pClientPort);
		m_pClientPort = NULL;
	}
	FltUnregisterFilter(m_pFilterHandle);
	m_pFilterHandle = NULL;
	return STATUS_SUCCESS;
};

FLT_PREOP_CALLBACK_STATUS FLTAPI FsPreCallback(
	__inout PFLT_CALLBACK_DATA Data,
	__in PCFLT_RELATED_OBJECTS FltObjects,
	__deref_out_opt PVOID *CompletionContext
)
{
	UNREFERENCED_PARAMETER(FltObjects);
	UNREFERENCED_PARAMETER(CompletionContext);

	PAGED_CODE();

	FLT_PREOP_CALLBACK_STATUS result = FLT_PREOP_SUCCESS_NO_CALLBACK;
	PFILE_OBJECT pFileObject = Data->Iopb->TargetFileObject;
	PFLT_FILE_NAME_INFORMATION pNameInfo = NULL;
	UNICODE_STRING ustrFilePath = { 0 };

	if (!pFileObject || FsRtlIsPagingFile(pFileObject))
		return result;

	if (m_EventList->IsCapturing())
	{
		ULONG ProcessId = FltGetRequestorProcessId(Data);

		if (Data->Iopb->MajorFunction == IRP_MJ_CREATE)
		{
			if (NT_SUCCESS(FltGetFileNameInformation(Data, FLT_FILE_NAME_NORMALIZED | FLT_FILE_NAME_QUERY_ALWAYS_ALLOW_CACHE_LOOKUP, &pNameInfo)))
			{
				svc_fs_create_file_data *data = (svc_fs_create_file_data *)
					ExAllocatePoolWithTag(PagedPool, sizeof(svc_fs_create_file_data), 'TXSB');

				if (data)
				{
					RtlZeroMemory(data, sizeof(svc_fs_create_file_data));

					RtlInitEmptyUnicodeString(&ustrFilePath, data->FilePath, sizeof(data->FilePath) - sizeof(WCHAR));

					data->protocol = svc_fs_create_file;
					data->size = sizeof(svc_fs_create_file_data);
					data->ProcessId = ProcessId;
					data->DesiredAccess = Data->Iopb->Parameters.Create.ShareAccess;
					data->Disposition = (Data->Iopb->Parameters.Create.Options >> 24) & 0x000000ff;
					data->Options = Data->Iopb->Parameters.Create.Options & 0x00ffffff;
					data->DesiredAccess = (Data->Iopb->Parameters.Create.SecurityContext) ? Data->Iopb->Parameters.Create.SecurityContext->DesiredAccess : 0;
					data->ShareAccess = Data->Iopb->Parameters.Create.ShareAccess;
					data->Attributes = Data->Iopb->Parameters.Create.FileAttributes;

					if (!NT_SUCCESS(NtFileNameToDosFileName(&pNameInfo->Name, &ustrFilePath)))
						RtlCopyUnicodeString(&ustrFilePath, &pNameInfo->Name);

					data->ResultStatus = STATUS_SUCCESS;//Modify later

					*CompletionContext = data;
					result = FLT_PREOP_SUCCESS_WITH_CALLBACK;
				}
				FltReleaseFileNameInformation(pNameInfo);
			}
		}
		else if (Data->Iopb->MajorFunction == IRP_MJ_CLEANUP)
		{
			svc_fs_close_file_data *data = (svc_fs_close_file_data *)
				ExAllocatePoolWithTag(PagedPool, sizeof(svc_fs_close_file_data), 'TXSB');

			if (data)
			{
				RtlZeroMemory(data, sizeof(svc_fs_close_file_data));

				RtlInitEmptyUnicodeString(&ustrFilePath, data->FilePath, sizeof(data->FilePath) - sizeof(WCHAR));

				data->protocol = svc_fs_close_file;
				data->size = sizeof(svc_fs_close_file_data);
				data->ProcessId = ProcessId;

				data->ResultStatus = STATUS_SUCCESS;//Modify later

				*CompletionContext = data;
				result = FLT_PREOP_SUCCESS_WITH_CALLBACK;
			}
		}
		else if (Data->Iopb->MajorFunction == IRP_MJ_READ)
		{
			svc_fs_readwrite_file_data *data = (svc_fs_readwrite_file_data *)
				ExAllocatePoolWithTag(PagedPool, sizeof(svc_fs_readwrite_file_data), 'TXSB');

			if (data)
			{
				RtlZeroMemory(data, sizeof(svc_fs_readwrite_file_data));

				RtlInitEmptyUnicodeString(&ustrFilePath, data->FilePath, sizeof(data->FilePath) - sizeof(WCHAR));

				data->protocol = svc_fs_readwrite_file;
				data->size = sizeof(svc_fs_readwrite_file_data);
				data->ProcessId = ProcessId;
				data->IsWrite = FALSE;
				data->ReadLength = Data->Iopb->Parameters.Read.Length;
				data->ByteOffset = Data->Iopb->Parameters.Read.ByteOffset.QuadPart;
				data->ResultStatus = STATUS_SUCCESS;//Modify later

				*CompletionContext = data;
				result = FLT_PREOP_SUCCESS_WITH_CALLBACK;
			}
		}
		else if (Data->Iopb->MajorFunction == IRP_MJ_WRITE)
		{
			svc_fs_readwrite_file_data *data = (svc_fs_readwrite_file_data *)
				ExAllocatePoolWithTag(PagedPool, sizeof(svc_fs_readwrite_file_data), 'TXSB');

			if (data)
			{
				RtlZeroMemory(data, sizeof(svc_fs_readwrite_file_data));

				RtlInitEmptyUnicodeString(&ustrFilePath, data->FilePath, sizeof(data->FilePath) - sizeof(WCHAR));

				data->protocol = svc_fs_readwrite_file;
				data->size = sizeof(svc_fs_readwrite_file_data);
				data->ProcessId = ProcessId;
				data->IsWrite = TRUE;
				data->ReadLength = Data->Iopb->Parameters.Write.Length;
				data->ByteOffset = Data->Iopb->Parameters.Write.ByteOffset.QuadPart;
				data->ResultStatus = STATUS_SUCCESS;//Modify later

				*CompletionContext = data;
				result = FLT_PREOP_SUCCESS_WITH_CALLBACK;
			}
		}
		//Get filename
		if (Data->Iopb->MajorFunction != IRP_MJ_CREATE && ustrFilePath.Buffer != NULL)
		{
			PFLT_CONTEXT pContext = NULL;
			if (NT_SUCCESS(FltGetStreamContext(FltObjects->Instance, pFileObject, &pContext)))
			{
				flt_file_context_t *ctx = (flt_file_context_t *)pContext;
				RtlCopyUnicodeString(&ustrFilePath, &ctx->DosFilePath);

				FltReleaseContext(pContext);
			}
		}
	}

	return result;
}

FLT_POSTOP_CALLBACK_STATUS FLTAPI FsPostCallbackSafe(
	_Inout_ PFLT_CALLBACK_DATA Data,
	_In_ PCFLT_RELATED_OBJECTS FltObjects,
	_In_opt_ PVOID CompletionContext,
	_In_ FLT_POST_OPERATION_FLAGS Flags
) {
	UNREFERENCED_PARAMETER(Flags);
	UNREFERENCED_PARAMETER(FltObjects);

	svc_nop_data *header = (svc_nop_data *)CompletionContext;

	if (header->protocol == svc_fs_create_file)
	{
		svc_fs_create_file_data *data = (svc_fs_create_file_data *)CompletionContext;
		data->ResultStatus = Data->IoStatus.Status;
	}
	else if (header->protocol == svc_fs_close_file)
	{
		svc_fs_close_file_data *data = (svc_fs_close_file_data *)CompletionContext;
		data->ResultStatus = Data->IoStatus.Status;
	}
	else if (header->protocol == svc_fs_readwrite_file)
	{
		svc_fs_readwrite_file_data *data = (svc_fs_readwrite_file_data *)CompletionContext;
		data->ResultStatus = Data->IoStatus.Status;
	}
	m_EventList->SendEvent(header);

	return FLT_POSTOP_FINISHED_PROCESSING;
}

FLT_POSTOP_CALLBACK_STATUS FLTAPI FsPostCallback(
	_Inout_ PFLT_CALLBACK_DATA Data,
	_In_ PCFLT_RELATED_OBJECTS FltObjects,
	_In_opt_ PVOID CompletionContext,
	_In_ FLT_POST_OPERATION_FLAGS Flags
) {
	UNREFERENCED_PARAMETER(FltObjects);
	UNREFERENCED_PARAMETER(Flags);
	UNREFERENCED_PARAMETER(CompletionContext);

	PFILE_OBJECT pFileObject = Data->Iopb->TargetFileObject;
	FLT_POSTOP_CALLBACK_STATUS result = FLT_POSTOP_FINISHED_PROCESSING;

	if (FlagOn(FLTFL_POST_OPERATION_DRAINING, Flags))
	{
		if(CompletionContext != NULL)
			ExFreePoolWithTag(CompletionContext, 'TXSB');

		return result;
	}

	//Save filename in context when post-create succeeded
	if (pFileObject && Data->Iopb->MajorFunction == IRP_MJ_CREATE && Data->IoStatus.Status == STATUS_SUCCESS)
	{
		//Make sure we have a stream context attached to FO
		PFLT_CONTEXT pContext = NULL;
		if (NT_SUCCESS(FltGetStreamContext(FltObjects->Instance, pFileObject, &pContext)))
		{
			FltReleaseContext(pContext);
		}
		else
		{
			PFLT_FILE_NAME_INFORMATION pNameInfo = NULL;
			NTSTATUS statusName = FltGetFileNameInformation(Data, FLT_FILE_NAME_NORMALIZED | FLT_FILE_NAME_QUERY_ALWAYS_ALLOW_CACHE_LOOKUP, &pNameInfo);
			if (!NT_SUCCESS(statusName))
			{
				statusName = FltGetFileNameInformation(Data, FLT_FILE_NAME_OPENED | FLT_FILE_NAME_QUERY_ALWAYS_ALLOW_CACHE_LOOKUP, &pNameInfo);
			}

			if (NT_SUCCESS(statusName))
			{
				if (NT_SUCCESS(FltAllocateContext(m_pFilterHandle, FLT_STREAM_CONTEXT, sizeof(flt_file_context_t), PagedPool, &pContext)))
				{
					if (NT_SUCCESS(FltSetStreamContext(FltObjects->Instance, pFileObject, FLT_SET_CONTEXT_REPLACE_IF_EXISTS, pContext, NULL)))
					{
						flt_file_context_t *ctx = (flt_file_context_t *)pContext;
						ctx->NtFilePath.Buffer = (PWCH)ExAllocatePoolWithTag(PagedPool, pNameInfo->Name.Length, 'TXSB');
						ctx->NtFilePath.Length = 0;
						ctx->NtFilePath.MaximumLength = pNameInfo->Name.Length;

						ctx->DosFilePath.Buffer = (PWCH)ExAllocatePoolWithTag(PagedPool, pNameInfo->Name.Length, 'TXSB');
						ctx->DosFilePath.Length = 0;
						ctx->DosFilePath.MaximumLength = pNameInfo->Name.Length;

						RtlCopyUnicodeString(&ctx->NtFilePath, &pNameInfo->Name);
						//This may fail... if fail, use NtFileName to replace it
						if (!NT_SUCCESS(NtFileNameToDosFileName(&pNameInfo->Name, &ctx->DosFilePath)))
							RtlCopyUnicodeString(&ctx->DosFilePath, &pNameInfo->Name);
					}
					//Decrese the reference
					FltReleaseContext(pContext);
				}

				FltReleaseFileNameInformation(pNameInfo);
			}
		}
	}

	if (CompletionContext != NULL)
	{
		if (KeGetCurrentIrql() < DISPATCH_LEVEL)
		{
			result = FsPostCallbackSafe(Data, FltObjects, CompletionContext, Flags);
		}
		else if (FLT_IS_IRP_OPERATION(Data))
		{
			FltDoCompletionProcessingWhenSafe(Data, FltObjects, CompletionContext, Flags, FsPostCallbackSafe, &result);
		}
		else
		{
			//this should never happen
			//DbgBreakPoint();
		}
	}

	return result;
}

const FLT_OPERATION_REGISTRATION Callbacks[] =
{
	{
		IRP_MJ_CREATE,
		0,
		FsPreCallback,
		FsPostCallback
	},
	{
		IRP_MJ_CLEANUP,
		0,
		FsPreCallback,
		FsPostCallback
	},
	{
		IRP_MJ_READ,
		FLTFL_OPERATION_REGISTRATION_SKIP_PAGING_IO,
		FsPreCallback,
		FsPostCallback
	},
	{
		IRP_MJ_WRITE,
		FLTFL_OPERATION_REGISTRATION_SKIP_PAGING_IO,
		FsPreCallback,
		FsPostCallback
	},
	/*{
		IRP_MJ_QUERY_INFORMATION,
		0,
		FsPreCallback,
		FsPostCallback
	},
	{
		IRP_MJ_DIRECTORY_CONTROL,
		0,
		FsPreCallback,
		FsPostCallback
	},*/
	{
		IRP_MJ_OPERATION_END
	}
};


VOID FLTAPI FsFileContextCleanup(
	_In_ PFLT_CONTEXT Context,
	_In_ FLT_CONTEXT_TYPE ContextType)
{
	UNREFERENCED_PARAMETER(ContextType);

	flt_file_context_t *ctx = (flt_file_context_t *)Context;
	ExFreePoolWithTag(ctx->NtFilePath.Buffer, 'TXSB');
	ExFreePoolWithTag(ctx->DosFilePath.Buffer, 'TXSB');
}

const FLT_CONTEXT_REGISTRATION ContextRegistration[] = {

	{ FLT_STREAM_CONTEXT,
	0,
	FsFileContextCleanup,
	sizeof(flt_file_context_t),
	'TXSB' },

	{ FLT_CONTEXT_END }
};

const FLT_REGISTRATION FilterRegistration =
{
	sizeof(FLT_REGISTRATION),         //  Size
	FLT_REGISTRATION_VERSION,           //  Version
	0,                                  //  Flags
	ContextRegistration,                               //  Context
	Callbacks,                          //  Operation callbacks
	FsUnload,                           //  MiniFilterUnload
	NULL,								//  InstanceSetup
	NULL,								//  InstanceQueryTeardown
	NULL,								//  InstanceTeardownStart
	NULL,								//  InstanceTeardownComplete
	NULL,                               //  GenerateFileName
	NULL,                               //  GenerateDestinationFileName
	NULL                                //  NormalizeNameComponent
};

NTSTATUS DriverEntryFilter(__in PDRIVER_OBJECT DriverObject, __in PUNICODE_STRING RegistryPath)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	PSECURITY_DESCRIPTOR sd;
	OBJECT_ATTRIBUTES oa;
	UNICODE_STRING ustrName;

	UNREFERENCED_PARAMETER(RegistryPath);
	__try
	{
		status = FltRegisterFilter(DriverObject, &FilterRegistration, &m_pFilterHandle);
		if (!NT_SUCCESS(status))
			__leave;

		status = FltBuildDefaultSecurityDescriptor(&sd, FLT_PORT_ALL_ACCESS);
		if (!NT_SUCCESS(status))
			__leave;

		RtlInitUnicodeString(&ustrName, L"\\SyscallMonPort");
		InitializeObjectAttributes(&oa, &ustrName, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, sd);
		status = FltCreateCommunicationPort(m_pFilterHandle, &m_pServerPort, &oa, NULL, FsConnectNotifyCallback, FsDisconnectNotifyCallback, FsMessageNotifyCallback, 1);
		FltFreeSecurityDescriptor(sd);

		status = FltStartFiltering(m_pFilterHandle);
		if (!NT_SUCCESS(status))
			__leave;

		PsCreateSystemThread(&m_EventList->m_hMsgThread, THREAD_ALL_ACCESS, NULL, NULL, NULL, MessageSenderThread, NULL);
	}
	__finally
	{
		if (!NT_SUCCESS(status))
		{
			if (NULL != m_pServerPort) {
				FltCloseCommunicationPort(m_pServerPort);
				m_pServerPort = NULL;
			}
			if (NULL != m_pFilterHandle) {
				FltUnregisterFilter(m_pFilterHandle);
				m_pFilterHandle = NULL;
			}
		}
	}
	return status;
}

VOID InitDynVers(VOID)
{
	m_IgnoreProcList = new CProcList();
	m_IgnoreFileList = new CFileList();
	m_EventList = new CEventList();
}

VOID FreeDynVers(VOID)
{
	delete m_IgnoreProcList;
	m_IgnoreProcList = NULL;

	delete m_IgnoreFileList;
	m_IgnoreFileList = NULL;

	delete m_EventList;
	m_EventList = NULL;
}

}//EXTERN C